import torch.nn as nn
import torch.optim as optim
from model.BiLSTM import *
from model.BiLSTM_CRF import *
from utils.data_loader import *
from tqdm import tqdm

# 实例化模型
models = {'BiLSTM': NERBiLSTM,
          'BiLSTM_CRF': NERBiLSTM_CRF}
model = models[config.model](config.embedding_dim, config.hidden_dim, config.dropout, word2id, config.tag2id)
model.load_state_dict(torch.load(f'save_model/{config.model}_best.pth',weights_only=False))

id2tag = {value: key for key, value in config.tag2id.items()}
# print(id2tag)

def model2predict(sample:str):
    """
    输入样本，直接提取出样本的实体
    :param sample: 样本，字符串类型
    :return:{实体：类型}，例如{'冠心病': 'DISEASE', '糖尿病': 'DISEASE'}
    """
    # 数值化
    x = []
    for char in sample:
        if char not in word2id:
            char = 'UNK'
        x.append(word2id[char])
    # 张量化
    x_tensor = torch.tensor([x])
    # 创建mask：模型输入需要
    mask = (x_tensor != 0).long()
    # 评估模型
    model.eval()
    with torch.no_grad():
        if model.name == "BiLSTM":
            outputs = model(x_tensor,mask)
            predicts = torch.argmax(outputs,dim=-1)[0]
            tags = [id2tag[i.item()] for i in predicts]
            print(tags)
        elif model.name == "BiLSTM_CRF":
            preds_ids = model(x_tensor,mask)[0]
            tags = [id2tag[i] for i in preds_ids]
        chars = [i for i in sample]
        assert len(chars) == len(tags)
        results = extract_entities(chars, tags)
        return results


def extract_entities(tokens, labels):
    """
    通过预测标签，从输入字符中，提取实体和类型
    :param tokens:输入字符，list，例如['冠','心','病']
    :param labels:预测标签，list,例如['B-DISEASE', 'I-DISEASE', 'I-DISEASE']
    :return:{实体：类型}，例如{'冠心病': 'DISEASE', '糖尿病': 'DISEASE'}
    """
    entities = []
    entity = []
    entity_type = None
    for token,label in zip(tokens,labels):
        if label.startswith('B-'):
            if entity:
                entities.append((entity_type, "".join(entity)))
                entity = []
                entity_type = None
            entity_type = label.split("-")[1]
            entity.append(token)
        elif label.startswith('I-') and entity:
            entity.append(token)
        else:
            if entity:
                entities.append((entity_type,"".join(entity)))
                entity = []
                entity_type = None

    if entity:
        entities.append((entity_type,''.join(entity)))
    return {entity: entity_type for entity_type, entity in entities}


if __name__ == '__main__':
    sample = '小明的父亲患有冠心病及糖尿病，无手术外伤史及药物过敏史糖尿病'
    result = model2predict(sample)
    print(result)